import pandas as pd
import re

def CalMass(formula):
    """
    Calculate the molecular mass from a chemical formula
    
    Args:
        formula (str): The molecular formula (e.g., 'C12H22O11')
    
    Returns:
        float: The calculated molecular mass
    """
    # Mass table for common elements (from original script)
    mass_table = {
        'C': 12.0000,
        'H': 1.007825,
        'O': 15.994914,
        'N': 14.003074,
        'S': 31.9721,
        'P': 30.9738,
        'Cl': 34.9689,
        'Na': 22.9898,
        'K': 38.9637,
        'Br': 78.9183,
        'F': 18.9984,
        'I': 126.9045,
        'D': 2.0141
    }
    
    # Count the atoms in the formula
    dic = {}
    for key in mass_table.keys():
        all_item_of_key = re.findall(key, formula)
        if all_item_of_key:
            pattern = re.compile(f'{key}[\d]+|{key}[a-z]+|{key}')
            num_str = re.findall(pattern, formula)
            num = 0
            for s in num_str:
                len_dic = s.split(key)
                if len(len_dic) > 1:  # Ensure there's something after the element symbol
                    if len_dic[1].isdigit():
                        num += int(len_dic[1])
                    elif len_dic[1] == '':
                        num += 1
                else:
                    num += 1  # Just the element with no number
            dic[key] = num
        else:
            dic[key] = 0
    
    # Calculate the total mass
    mass = 0
    for key in mass_table.keys():
        mass += (mass_table[key] * dic[key])
    return mass

def add_MolecularWeight(row, adjustment=0):
    """
    Calculate molecular weight based on formula with optional adjustment
    
    Args:
        row: DataFrame row containing 分子式 column
        adjustment (float): Value to add to the calculated mass
    
    Returns:
        str: Formatted molecular weight or empty string if no formula
    """
    if pd.isna(row['分子式']) or row['分子式'] == "":
        return ""
    else:
        try:
            mass = CalMass(row['分子式']) + adjustment
            return format(mass, '.4f')
        except Exception as e:
            print(f"Error calculating mass for {row['分子式']}: {e}")
            return ""

# Ion adjustments for the specific columns you want to calculate
targeted_adjustments = {
    'M_PLUS_H_PLUS': 1.007825,  # [M+H]+
    'M_PLUS_Na_PLUS': 22.9898,  # [M+Na]+
    'M_PLUS_K_PLUS': 38.9637,  # [M+K]+
    'M_PLUS_NH4_PLUS': 18.034374,  # [M+NH4]+
    'M_MIUNS_H_MINUS': -1.007825,  # [M-H]-
    'M_PLUS_COOH_MINUS': 44.997653,  # [M+FA-H]- (Formic Acid)
}

def main():
    # File paths
    input_file = './Data/xbhb_new.xlsx'
    output_file = './Data/xbhb_new_updated.xlsx'
    
    try:
        # Load the Excel file
        print(f"Loading file: {input_file}")
        df = pd.read_excel(input_file)
        print(f"Loaded file with {len(df)} rows and {len(df.columns)} columns")
        
        # Display first few rows to understand the structure
        print("\nFirst 3 rows of data:")
        print(df.head(3))
        
        # Display column names
        print("\nColumn names:")
        for i, col in enumerate(df.columns):
            print(f"{i}: {col}")
        
        # Check if 分子式 column exists
        if '分子式' not in df.columns:
            print("\nWarning: '分子式' column not found!")
            formula_col = input("Enter the name of the column containing molecular formulas: ")
            if formula_col.strip() and formula_col in df.columns:
                df.rename(columns={formula_col: '分子式'}, inplace=True)
            else:
                print("Cannot proceed without molecular formula data.")
                return
        
        # Preview some molecular formulas
        print("\nSample molecular formulas:")
        formulas = df['分子式'].dropna().head(5).tolist()
        for i, formula in enumerate(formulas):
            print(f"{i+1}: {formula}")
        
        # Verify with first row before proceeding
        print("\nVerifying calculations using first row...")
        first_row = df.iloc[0]
        
        for col, adjustment in targeted_adjustments.items():
            if col in df.columns and pd.notna(first_row[col]) and first_row[col] != "":
                calculated = add_MolecularWeight(first_row, adjustment)
                existing = str(first_row[col])
                
                print(f"Column: {col}")
                print(f"  Formula: {first_row['分子式']}")
                print(f"  Original value: {existing}")
                print(f"  Calculated value: {calculated}")
                
                # Check if values match within rounding tolerance
                try:
                    diff = abs(float(calculated) - float(existing))
                    if diff < 0.0001:
                        print(f"  ✓ Values match (difference: {diff:.6f})")
                    else:
                        print(f"  ✗ Values differ by: {diff:.6f}")
                        proceed = input(f"Values for {col} don't match. Continue? (y/n): ")
                        if proceed.lower() != 'y':
                            print("Aborting calculation.")
                            return
                except ValueError:
                    print("  ✗ Couldn't compare numeric values")
                    proceed = input(f"Unable to verify {col}. Continue? (y/n): ")
                    if proceed.lower() != 'y':
                        print("Aborting calculation.")
                        return
        
        # Calculate the specific columns
        print("\nCalculating ion masses...")
        # Create a copy of the dataframe to preserve original values in first row
        df_new = df.copy()
        
        # Skip first row for calculation (use it as verification only)
        for col, adjustment in targeted_adjustments.items():
            if col not in df_new.columns:
                df_new[col] = ""
            
            # Apply calculation to all rows except the first
            for i in range(1, len(df_new)):
                row = df_new.iloc[i]
                df_new.at[i, col] = add_MolecularWeight(row, adjustment)
            
            print(f"Calculated {col} for {len(df_new)-1} rows (skipped first row)")
        
        # Save the results
        df_new.to_excel(output_file, index=False)
        print(f"\nResults saved to {output_file}")
        
        # Preview results
        print("\nSample results (first 3 rows):")
        columns_to_show = ['分子式'] + list(targeted_adjustments.keys())
        print(df_new[columns_to_show].head(3))
        
    except Exception as e:
        print(f"\nAn error occurred: {e}")

if __name__ == "__main__":
    main()